# -*- coding: utf-8 -*-
# time: 2025/5/15 09:06
# file: ch01.py
# author: hanson
"""
sk-62fb03f9198f4720bb359c873757fafc
https://bailian.console.aliyun.com/?tab=app#/api-key
"""
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain.tools import Tool
from langchain_community.utilities.sql_database import SQLDatabase
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.chat_models.tongyi import ChatTongyi
from langchain.agents import initialize_agent, AgentType
from langchain_ollama import ChatOllama

# 1. 初始化Qwen大模型 ： 使用ChatTongyi封装类连接Dashscope的Qwen模型
llm = ChatTongyi(
    model_name="qwen-max",
    dashscope_api_key="sk-62fb03f9198f4720bb359c873757fafc",
    temperature=0.3  # 降低随机性
)
# llm = ChatOllama(model="qwen2.5:1.5b", temperature=0.3)  #本地模型太差了
# 2. 连接MySQL数据库
db = SQLDatabase.from_uri(
        "mysql+pymysql://root:1234qwer@192.168.6.43:3306/iot_ts",
    sample_rows_in_table_info=3
)
"""
# 在SQLDatabase初始化时添加安全限制
db = SQLDatabase.from_uri(
    uri,
    include_tables=['allowed_table1', 'allowed_table2'],  # 白名单
    sample_rows_in_table_info=3,
    max_string_length=100  # 限制返回字符串长度
)

使用SQL Agent Toolkit（更专业的SQL处理）

from langchain_community.agent_toolkits import create_sql_agent

sql_agent = create_sql_agent(
    llm=llm,
    db=db,
    agent_type="openai-tools",
    verbose=True
)
"""

# 3. 创建数据库查询工具
def query_database(query: str) -> str:
    """执行SQL查询并返回结果"""
    try:
        result = db.run(query)
        return f"查询结果: {str(result)}"
    except Exception as e:
        return f"查询错误: {str(e)}"

# 创建数据库查询工具：工具描述会影响Agent何时选择使用它
db_tool = Tool(
    name="database_query",
    func=query_database,
    description="用于查询MySQL数据库。输入必须是有效的SQL语句。"
)


# 4. 创建加法工具
def add_numbers(numbers: str) -> str:
    """将两个数字相加，输入格式如 '3,5'"""
    try:
        a, b = map(float, numbers.split(","))
        return str(a + b)
    except Exception as e:
        return f"计算错误: {str(e)}"

# 创建数据库查询工具：工具描述会影响Agent何时选择使用它
add_tool = Tool(
    name="add_numbers",
    func=add_numbers,
    description="用于两个数字相加。输入必须是逗号分隔的两个数字，如 '3,5'。"
)

# 5. 使用更兼容的Agent初始化方式
tools = [db_tool, add_tool]

"""
使用 ZERO_SHOT_REACT_DESCRIPTION 替代原来的 create_tool_calling_agent
这是目前对Qwen兼容性最好的Agent类型
"""
# 新版推荐使用ZERO_SHOT_REACT_DESCRIPTION类型，详细执行日志（verbose=True）
agent = initialize_agent(
    tools,
    llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True,
    handle_parsing_errors=True  # 关键：处理解析错误
)

# 6. 测试运行
if __name__ == "__main__":
    queries = [
        "users表有多少条记录？",
        "计算123.5加456.7等于多少",
        "先计算10加20，然后查询tk_device表的前3条记录"
    ]

    for query in queries:
        print(f"\n=== 执行查询: {query} ===")
        try:
            result = agent.run(query)
            print(f"结果: {result}")
        except Exception as e:
            print(f"执行出错: {str(e)}")